from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from torch.utils.data import DataLoader, Dataset

# 定义数据集类
class TextClassificationDataset(Dataset):
    """
    初始化数据处理类。

    参数:
    texts (list of str): 输入文本列表，每个元素为一个文本字符串。
    labels (list): 标签列表， 每个元素对应于texts列表中文本的标签。
    tokenizer: 用于文本分词的Tokenizer对象。
    max_length (int): 文本分词后的最大长度。
    """
    # 存储输入参数为实例变量
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        
        """
        根据索引idx返回文本数据及其对应的标签。

        参数:
        - idx: int ， 数据项的索引。

        返回:
        - 一个字典， 包含输入ID、注意力掩码和标签。
        """
        # 获取索引idx对应的文本和标签    
        text = self.texts[idx]
        label = self.labels[idx]
        # 使用tokenizer对文本进行编码
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,        # 添加特殊标记，如[BOS], [EOS]
            max_length=self.max_length,     # 设置最大序列长度
            padding='max_length',           # 使用最大长度进行填充
            truncation=True,                # 超过最大长度时进行截断
            return_tensors='pt'             # 返回PyTorch张量
        )
        # 返回编码后的文本数据和对应的标签
        return {
            'input_ids': encoding['input_ids'].flatten(),               # 输入ID
            'attention_mask': encoding['attention_mask'].flatten(),     # 注意力掩码
            'labels': torch.tensor(label, dtype=torch.long)             # 标签
        }

# 示例数据
texts = ["这是一个正向评价语句。", "这是一个负向评价语句。"]
labels = [1, 0]     # 1表示积极，0表示消极

# 加载模型和分词器
model_name = r'D:\VS_workspace\A002_AIPython学习\Models_cache\transformers_roberta_base_chinese_cluecorpussmall'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 创建数据集和数据加载器
dataset = TextClassificationDataset(texts, labels, tokenizer, max_length=128)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# 定义优化器和损失函数
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
criterion = torch.nn.CrossEntropyLoss()

# 微调模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

for epoch in range(3):
    for batch in dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask)
        logits = outputs.logits
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch + 1} completed.')


"""
在这个示例中， 我们定义了一个简单的文本分类数据集， 使用AutoModelForSequenceClassification加载预训练的 RoBERTa 模型。
然后，借助 PyTorch 的DataLoader批量加载数据 ， 使用 AdamW 优化器和交叉熵损失函数对模型进行微调。
通过以上方式，你就可以获取和使用 RoBERTa 模型进行自然语言处理任务。

这段代码中：

使用的是 Hugging Face 提供的 AutoModelForSequenceClassification 模型，基于 RoBERTa 的中文预训练模型。
通过加载本地路径下的模型权重文件 transformers_roberta_base_chinese_cluecorpussmall 获取了一个已经完成预训练的模型。
在此基础上， 对模型进行微调 （ Fine-tuning ） ， 使其能够更好地处理我们自定义的文本分类任务。

微调的目的可以总结为以下几点：
迁移学习：利用预训练模型在大规模语料库上学习到的语言表示能力，迁移到当前的小规模、特定任务数据集上。
任务适配： RoBERTa 预训练模型本身是通用语言模型，不具备直接做文本分类的能力。通过微调，可以将其适配到具体的文本分类任务。
参数更新：在训练过程中，优化器会根据损失函数（如交叉熵）反向传播并更新模型参数，使模型输出更符合当前任务的真实标签。
提高准确率：相比从头开始训练模型，微调通常能更快收敛，并在有限的数据量下获得更好的泛化能力。
在你提供的代码中，微调过程包括以下几个关键步骤：

将模型移动到 GPU 或 CPU 设备上。
定义优化器 （ AdamW ） 和损失函数 （ CrossEntropyLoss ） 。
进行多轮训练 （epoch ），每次从数据加载器中取出一个 batch 的数据，前向传播得到预测结果，计算损失后进行反向传播和参数更新。
打印每个 epoch 的训练完成信息。
这样经过若干轮训练后，模型就能够在你的特定任务（例如中文情感分类）上取得较好的效果。

"""